# 60年代风格图形渲染器（向量化版本）
# 使用AI辅助

import numpy as np
import scipy
import scipy.io
import PIL.Image

import multiprocessing as mp
from multiprocessing import Queue, Process
import time

import numba
import typing

is_pbr_ready = False

try:
    import drawsphere_pbr
    #raise ValueError('123')
    is_pbr_ready = True
except Exception as e:
    print(f"Failed to Import PBR Module: {e}. Use Blinn-Phong instead. ")

allow_normal_map = False

class ADVMAT:
    def __init__(self, F0 = None, metallic = 0.0, roughness=0.5):
        self.F0 = 0.04*np.ones((3))
        self.metallic = float(metallic)
        self.roughness = float(roughness)

        if F0 is not None:
            self.F0[:] = F0
        return

class GROUNDS:
    def __init__(self, ground_y, advmat = None):
        self.ground_y = ground_y

        self.advmat = advmat if advmat else ADVMAT()
        return

    def is_ray_collided_with_me(self, points, directions):
        # points: (N,3), directions: (N,3)
        N = points.shape[0]
        norms = np.linalg.norm(directions, axis=1, keepdims=True)
        directions = directions / norms

        is_collided = np.zeros(N, dtype=bool)
        collided_coord = np.zeros((N, 3))
        normal = np.zeros((N, 3))
        t = np.zeros(N)
        t_final = np.zeros(N)
        color = np.zeros((N, 3))

        # 向量化计算
        valid_mask = ~np.isclose(directions[:, 1], 0.0)
        valid_indices = np.where(valid_mask)[0]
        t[valid_indices] = (self.ground_y - points[valid_indices, 1]) / directions[valid_indices, 1]

        valid_mask = valid_mask & (t > 0)
        valid_indices = np.where(valid_mask)[0]

        if valid_indices.size > 0:
            is_collided[valid_indices] = True
            collided_coord[valid_indices] = points[valid_indices] + t[valid_indices, None] * directions[valid_indices]
            t_final[valid_indices] = t[valid_indices]

            # 向量化获取法向量和颜色
            normal[valid_indices] = self.get_norm(len(valid_indices))
            color[valid_indices] = self.get_color(collided_coord[valid_indices])

        return is_collided, collided_coord, normal, color, t_final

    def get_norm(self, n=1):
        normal = np.tile([0.0, 1.0, 0.0], (n, 1))
        if allow_normal_map:
            perturbation = 0.05 * (2 * np.random.rand(n, 3) - 1)
            normal += perturbation
        norms = np.linalg.norm(normal, axis=1, keepdims=True)
        normal = normal / norms
        return normal

    def get_color(self, coords):
        theta = np.pi * 1 / 6
        rot = np.array([[np.cos(theta), 0, np.sin(theta)],
                        [0, 1, 0],
                        [-np.sin(theta), 0, np.cos(theta)]])
        _coords = (rot @ coords.T).T

        a = np.floor(_coords[:, 0] % 2)
        b = np.floor(_coords[:, 2] % 2)

        condition = (a == 0) & (b == 0) | (a == 1) & (b == 1)
        colors = np.where(condition[:, None], [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
        return colors

class SPHERES:
    def __init__(self, coord, r, color=None,advmat=None):
        self.coord = np.array(coord)
        self.color = np.ones(3) if color is None else np.array(color)
        self.r = float(r)
        self.box_size = np.array([2*r,2*r,2*r])
        self.advmat = advmat if advmat else ADVMAT()
        return

    def is_ray_collided_with_me(self, points, directions):
        if points.ndim != 2 or directions.ndim !=2:
            raise ValueError('123')
            
        N = points.shape[0]
        norms = np.linalg.norm(directions, axis=1, keepdims=True)
        directions = directions / norms

        is_collided = np.full(N, False, dtype=bool)
        collided_coord = np.empty((N, 3))
        normal = np.empty((N, 3))
        color = np.tile(self.color, (N, 1))
        t_final = np.full(N, np.inf)
        
        is_collided, t_final = self._is_ray_collided_with_me_numba(points, directions, self.coord, self.r)
        
        collided_mask = np.where(is_collided)[0]
        collided_coord[collided_mask] = points[collided_mask] + t_final[collided_mask, None] * directions[collided_mask]
        normal[collided_mask] = self.get_norm(collided_coord[collided_mask])
        return is_collided, collided_coord, normal, color, t_final

    @staticmethod
    @numba.njit(fastmath=True, cache=True)
    def _is_ray_collided_with_me_numba(points, directions, coord, r):
        N = points.shape[0]
        
        is_collided = np.zeros(N, dtype=np.bool_)
        t_final = np.full(N, np.inf)
        
        for i in range(N):
            # Calculate difference vector
            diff_x = points[i,0] - coord[0]
            diff_y = points[i,1] - coord[1]
            diff_z = points[i,2] - coord[2]
            
            # Calculate coefficients
            A = (directions[i,0]**2 + directions[i,1]**2 + directions[i,2]**2)
            B = 2 * (directions[i,0]*diff_x + directions[i,1]*diff_y + directions[i,2]*diff_z)
            C = (diff_x**2 + diff_y**2 + diff_z**2) - r**2
            
            # Calculate discriminant
            D = B**2 - 4*A*C
            
            if D > 0:
                sqrt_D = np.sqrt(D)
                t1 = (-B + sqrt_D) / (2*A)
                t2 = (-B - sqrt_D) / (2*A)
                
                # Find smallest positive t
                min_t = np.inf
                if t1 > 0 and t1 < min_t:
                    min_t = t1
                if t2 > 0 and t2 < min_t:
                    min_t = t2
                    
                if min_t < np.inf:
                    is_collided[i] = True
                    t_final[i] = min_t
                    
        return is_collided, t_final
    
    def get_norm(self, points):
        normals = points - self.coord
        if allow_normal_map:
            perturbation = 0.05 * (2 * np.random.rand(points.shape[0], 3) - 1)
            normals += perturbation
        norms = np.linalg.norm(normals, axis=1, keepdims=True)
        normals = normals / norms
        return normals

class CUBES:
    def __init__(self, coord, lwh, color=None, advmat = None):
        self.coord = np.array(coord)
        self.lwh = np.array(lwh)
        self.color = np.ones(3) if color is None else np.array(color)

        self.advmat = advmat if advmat else ADVMAT()
        return

    def is_ray_collided_with_me(self, points, directions):
        if points.ndim != 2 or directions.ndim !=2:
            raise ValueError('123')
        """
        使用is_collided_with_AABB函数重构后的版本
        """
        N = points.shape[0]
        directions = directions / np.linalg.norm(directions, axis=1, keepdims=True)
        
        # 调用重构后的函数
        is_collided, t_min, t_max = is_collided_with_AABB_numba(points, directions, self.coord, self.lwh)
        
        # 初始化输出数组
        collided_coord = np.empty((N, 3))
        t_final = np.full(N, np.inf)
        normal = np.empty((N, 3))
        color = np.tile(self.color, (N, 1))
        
        valid_indices = np.where(is_collided)[0]
        if np.any(valid_indices):
            # 计算最终相交参数
            t_final[valid_indices] = np.where(
                t_min[valid_indices] > 0, 
                t_min[valid_indices], 
                t_max[valid_indices]
            )
            # 计算碰撞坐标
            collided_coord[valid_indices] = points[valid_indices] + t_final[valid_indices, None] * directions[valid_indices]
            # 计算法线
            normal[valid_indices] = self.get_norm(collided_coord[valid_indices])
        return is_collided, collided_coord, normal, color, t_final
    
    def get_norm(self, coords):
        if coords.ndim != 2:
            raise ValueError('123')
            
        local_coords = coords - self.coord[None,:]
        half_extents = self.lwh / 2.0
        normalized_coords = local_coords / half_extents
    
        axis = np.argmax(np.abs(normalized_coords), axis=1)
        signs = np.sign(normalized_coords[np.arange(len(axis)), axis])
    
        normals = np.zeros((len(axis), 3))
        normals[np.arange(len(axis)), axis] = signs
    
        if allow_normal_map:
            perturbation = 0.05 * (2 * np.random.rand(len(axis), 3) - 1)
            normals += perturbation
    
        norms = np.linalg.norm(normals, axis=1, keepdims=True)
        normals = normals / norms
        return normals

class SPRITERS:
    def __init__(self, coord, height, texture, advmat = None):
        self.coord = np.zeros((3))
        self.coord[:] = np.array(coord)

        self.texture = np.flip(np.array(texture).transpose([1, 0, 2]), axis=1)
        if self.texture.ndim != 3 or self.texture.shape[2] != 3:
            raise ValueError('Texture must be RGB image')

        self.height = float(height)
        self.scale = float(height) / self.texture.shape[1]

        self.advmat = advmat if advmat else ADVMAT()
        return

    def is_ray_collided_with_me(self, points, directions):
        if points.ndim != 2 or directions.ndim !=2:
            raise ValueError('123')
        
        N = points.shape[0]
        directions = directions / np.linalg.norm(directions,axis=1,keepdims=True)
        normal = np.tile((0,0,1),(N,1))
        is_collided, collided_coord, t_final, color = self._is_collided_numba(points,directions,self.coord, self.texture,self.scale)
        return is_collided, collided_coord, normal, color, t_final

    @staticmethod
    @numba.njit(fastmath=True,cache=True)
    def _is_collided_numba(points, directions, coord, texture, scale):
        N = points.shape[0]
        is_collided = np.full(N,False)
        collided_coord = np.empty((N,3))
        colors = np.empty((N, 3))
        t = np.full(N,np.inf)
        
        for k in range(N):
            if np.abs(directions[k, 2]) > 0:
                t[k] = (coord[2] - points[k, 2]) / directions[k, 2]
                is_collided[k] =  (t[k] > 0) & (~np.isinf(t[k]))
                collided_coord[k, :] = points[k, :] + t[k] * directions[k,:]
        
        local_xy = np.empty((3))
        for k in range(N):
            if is_collided[k]:
                local_xy = collided_coord[k,:] - coord
                i = int(local_xy[0] / scale + texture.shape[0] // 2)
                j = int(local_xy[1] / scale + texture.shape[1] // 2)
                if (i >= 0) & (i < texture.shape[0]) & (j >= 0) & (j < texture.shape[1]):
                    colors[k, :] = texture[i, j, :]
                    if np.sum(colors[k,:]**2) < 0.1:
                        is_collided[k] = False
                else:
                    is_collided[k] = False
        return is_collided, collided_coord, t, colors
        

class TRIANGLES:
    def __init__(self, coord, model = None, colors = None, advmat = None):
        self.triangles = np.zeros((8,3,3)) # 面-3顶点-xyz
        self.normals = None
        self.boundingbox = None
        self.coord = None
        self.box_size = None
        self.colors = np.ones((8,3))
        self.advmat = advmat if advmat else ADVMAT()

        # 定义6个顶点
        v0 = np.array([0, 0,  1 - 5])   # 顶部
        v1 = np.array([1, 0,  0 - 5])   # 右
        v2 = np.array([0, 1,  0 - 5])   # 前
        v3 = np.array([-1, 0, 0 - 5])  # 左
        v4 = np.array([0, -1, 0 - 5])  # 后
        v5 = np.array([0, 0, -1 - 5])  # 底部

        # 定义8个三角形面
        self.triangles[0,:,:] = np.array([v0, v1, v2])  # 顶-右-前
        self.triangles[1,:,:] = np.array([v0, v2, v3])  # 顶-前-左
        self.triangles[2,:,:] = np.array([v0, v3, v4])  # 顶-左-后
        self.triangles[3,:,:] = np.array([v0, v4, v1])  # 顶-后-右
        self.triangles[4,:,:] = np.array([v5, v2, v1])  # 底-前-右
        self.triangles[5,:,:] = np.array([v5, v3, v2])  # 底-左-前
        self.triangles[6,:,:] = np.array([v5, v4, v3])  # 底-后-左
        self.triangles[7,:,:] = np.array([v5, v1, v4])  # 底-右-后
        #self.triangles.flags.writeable = False
        self.init_shape()

        self.load(model, colors)
        self.moveto(coord)
        return

    def calculate_bounding_box(self):
        """
        计算多面体的包围盒

        返回:
            center (np.array): 包围盒的中心坐标 [x, y, z]
            dimensions (np.array): 包围盒的长宽高 [width, height, depth]
        """
        # 将所有顶点重塑为一个二维数组，每行是一个顶点的xyz坐标
        vertices = self.triangles.reshape(-1, 3)

        # 计算每个坐标轴的最小值和最大值
        min_coords = np.min(vertices, axis=0) - 0.1
        max_coords = np.max(vertices, axis=0) + 0.1

        # 计算包围盒的中心
        box_center = (min_coords + max_coords) / 2

        # 计算包围盒的长宽高
        box_size = max_coords - min_coords

        return box_center.flatten(), box_size.flatten()

    def init_shape(self):
        self.v0 = self.triangles[:,0,:] #(Ntri,3)
        self.v1 = self.triangles[:,1,:]
        self.v2 = self.triangles[:,2,:]
        self.edge1 = self.v1 - self.v0
        self.edge2 = self.v2 - self.v0

        self.normals = np.cross(self.edge1,self.edge2)
        self.normals /= np.linalg.norm(self.normals,axis=1,keepdims=True)

        self.coord, self.box_size = self.calculate_bounding_box()
        return

    def load(self, model = None, colors = None):
        '''
        读取类obj格式
        输入：N*3*3 矩阵，N是三角形个数(axis 0)，每个三角形3个顶点(axis 1)，和它的xyz坐标(axis 2)
        '''
        if model is None and colors is None:
            return

        if model.ndim == 3 and model.shape[1] == 3 and model.shape[2] == 3:
            N = model.shape[0]
            self.triangles = model.copy()
            self.init_shape()
            self.moveto(np.array([0,0,-10]))
            print(f'size of obj:{self.box_size}')
        else:
            raise ValueError('failed to load')

        if colors is None:
            self.colors = np.ones((N, 3))
        elif colors.ndim == 1 and colors.shape[0] == 3:
            self.colors = np.tile(colors,(N, 1))
        elif colors.ndim == 2 and colors.shape[0] == N and colors.shape[1] == 3:
            self.colors = colors.copy()
        else:
            raise ValueError('failed to load')

        return

    def moveto(self, point):
        if point.ndim != 1:
            raise ValueError('123')

        direction = point - self.coord
        self.move(direction)
        return

    def move(self, direction):
        if direction.ndim != 1:
            raise ValueError('123')

        self.triangles += direction[None,None,:]
        self.init_shape()
        return

    def rotate(self, axis, theta):
        R = get_rotation_matrix(axis,theta)
        triangles = self.triangles - self.coord[None,None,:] #(N,3,3) - (1,1,3)
        triangles = np.einsum('jk,ilk->ilj',R, triangles) #(3,3), (N,3,3) -> (N,3,3)
        triangles += self.coord[None,None,:]
        self.triangles = triangles
        self.init_shape()
        return

    def scale(self, scale):
        triangles = self.triangles - self.coord[None,None,:]
        triangles *= scale
        self.triangles = triangles + self.coord[None,None,:]
        self.init_shape()
        return

    def is_ray_collided_with_me(self, points, directions):
        '''
        带有AABB的版本
        '''
        if points.ndim != 2 or directions.ndim != 2:
            raise ValueError('123')

        N = points.shape[0]
        is_collided, collided_coord, normals, colors, t_final = np.full(N,False), np.empty((N,3)),np.empty((N,3)),np.empty((N,3)),np.empty((N))

        aabb_mask = is_collided_with_AABB_numba(points, directions, self.coord, self.box_size)[0]
        if np.any(aabb_mask):
            _is_collided, _collided_coord, _normals, _colors, _t_final = self._is_ray_collided_with_me(points[aabb_mask], directions[aabb_mask])
            
            is_collided[aabb_mask] = _is_collided
            collided_coord[aabb_mask] = _collided_coord
            normals[aabb_mask] = _normals
            colors[aabb_mask] = _colors
            t_final[aabb_mask] = _t_final

        return is_collided, collided_coord, normals, colors, t_final 

    def _is_ray_collided_with_me(self, points, directions):
        '''
        这个是真正的内部实现
        外部调用请使用上述带有AABB版本
        '''
        if points.ndim != 2 or directions.ndim != 2:
            raise ValueError('123')

        Npoints = points.shape[0]
        Ntri = self.triangles.shape[0]
        
        points = points[:,None,:] #(N_point,1,3)
        directions = directions[:,None,:]

        v0 = self.v0[None,:,:] #(1, N_tri,3)
        v1 = self.v1[None,:,:] #(1, N_tri,3)
        v2 = self.v2[None,:,:] #(1, N_tri,3)
        edge1 = self.edge1[None,:,:]
        edge2 = self.edge2[None,:,:]

        is_collided, t_final, t_min_indices = self._triangles_collided_numba(points, directions, v0, edge1, edge2)

        collided_coord = points[:,0,:] + t_final[:,None] * directions[:,0,:]
        normals = self.normals[t_min_indices].copy()
        colors = self.colors[t_min_indices].copy()

        isout = np.sum(normals * directions[:,0,:], axis = 1) > 0 # 选取正确的法向量（假设三角形围成封闭图形，总在外侧观察）
        normals[isout] *= -1
        #normals = np.tile([0,0,1],(Npoints,1))
        return is_collided, collided_coord, normals, colors, t_final

    @staticmethod
    def _collided(points, directions, v0, edge1, edge2):
        '''
        实际使用的是外部jit的版本
        '''
        h = np.cross(directions, edge2) #(N_point,1,3) (1,N_tri,3) -> (N_point,N_tri,3)
        det = np.sum(edge1 * h, axis = 2) #(N_point,Ntri)

        is_collided_global = np.abs(det) > 1e-3 # 排除平行的三角形
        det[~is_collided_global] = 1 # 因为不相交所以没有意义，避免除以0

        inv_det = 1 / det

        s = points - v0
        u = inv_det * np.sum(s * h, axis = 2)

        q = np.cross(s, edge1)
        v = inv_det * np.sum(directions * q, axis = 2)

        is_collided_global &= (u > 0) & (u < 1) & (v > 0) & (u + v < 1) # 排除不在范围内的三角形

        t = inv_det * np.sum(edge2 * q, axis = 2)
        t[(t < 0) | (~is_collided_global)] = np.inf

        t_min_indices = np.argmin(t,axis=1)
        t_final = np.min(t, axis=1)
        is_collided = ~np.isinf(t_final)
        return is_collided, t_final, t_min_indices

    @staticmethod
    @numba.njit(fastmath=True, cache=True)
    def _triangles_collided_numba(
        points: np.ndarray,          # shape (N_point, 1, 3)
        directions: np.ndarray,      # shape (N_point, 1, 3)
        v0: np.ndarray,              # shape (1, N_tri, 3)
        edge1: np.ndarray,           # shape (1, N_tri, 3)
        edge2: np.ndarray            # shape (1, N_tri, 3)
    ) -> typing.Tuple[np.ndarray, np.ndarray, np.ndarray]:
        N_point = points.shape[0]
        N_tri = v0.shape[1]

        # 预分配结果数组
        h = np.empty((N_point, N_tri, 3))
        det = np.empty((N_point, N_tri))
        is_collided_global = np.empty((N_point, N_tri), dtype=np.bool_)
        inv_det = np.empty((N_point, N_tri))
        u = np.empty((N_point, N_tri))
        v = np.empty((N_point, N_tri))
        t = np.empty((N_point, N_tri))

        # 计算h和det
        for i in range(N_point):
            for j in range(N_tri):
                # 计算叉积 h = directions × edge2
                h[i,j,0] = directions[i,0,1] * edge2[0,j,2] - directions[i,0,2] * edge2[0,j,1]
                h[i,j,1] = directions[i,0,2] * edge2[0,j,0] - directions[i,0,0] * edge2[0,j,2]
                h[i,j,2] = directions[i,0,0] * edge2[0,j,1] - directions[i,0,1] * edge2[0,j,0]

                # 计算点积 det = edge1 · h
                det[i,j] = edge1[0,j,0] * h[i,j,0] + edge1[0,j,1] * h[i,j,1] + edge1[0,j,2] * h[i,j,2]

                # 判断是否碰撞
                is_collided_global[i,j] = np.abs(det[i,j]) > 1e-3
                inv_det[i,j] = 1.0 / det[i,j] if is_collided_global[i,j] else 1.0

        # 计算u
        for i in range(N_point):
            for j in range(N_tri):
                s0 = points[i,0,0] - v0[0,j,0]
                s1 = points[i,0,1] - v0[0,j,1]
                s2 = points[i,0,2] - v0[0,j,2]

                dot_s_h = s0 * h[i,j,0] + s1 * h[i,j,1] + s2 * h[i,j,2]
                u[i,j] = inv_det[i,j] * dot_s_h

        # 计算v和t
        for i in range(N_point):
            for j in range(N_tri):
                s0 = points[i,0,0] - v0[0,j,0]
                s1 = points[i,0,1] - v0[0,j,1]
                s2 = points[i,0,2] - v0[0,j,2]

                # 计算叉积 q = s × edge1
                q0 = s1 * edge1[0,j,2] - s2 * edge1[0,j,1]
                q1 = s2 * edge1[0,j,0] - s0 * edge1[0,j,2]
                q2 = s0 * edge1[0,j,1] - s1 * edge1[0,j,0]

                # 计算点积 directions · q
                dot_dir_q = directions[i,0,0] * q0 + directions[i,0,1] * q1 + directions[i,0,2] * q2
                v[i,j] = inv_det[i,j] * dot_dir_q

                # 计算点积 edge2 · q
                dot_edge2_q = edge2[0,j,0] * q0 + edge2[0,j,1] * q1 + edge2[0,j,2] * q2
                t[i,j] = inv_det[i,j] * dot_edge2_q

                # 更新碰撞条件
                is_collided_global[i,j] &= (u[i,j] > 0) & (u[i,j] < 1) & (v[i,j] > 0) & (u[i,j] + v[i,j] < 1)

                if (t[i,j] < 0) or (not is_collided_global[i,j]):
                    t[i,j] = np.inf

        # 计算最终结果
        t_final = np.empty(N_point)
        t_min_indices = np.empty(N_point, dtype=np.int64)
        is_collided = np.empty(N_point, dtype=np.bool_)

        for i in range(N_point):
            min_t = np.inf
            min_idx = 0
            for j in range(N_tri):
                if t[i,j] < min_t:
                    min_t = t[i,j]
                    min_idx = j

            t_final[i] = min_t
            t_min_indices[i] = min_idx
            is_collided[i] = not np.isinf(min_t)

        return is_collided, t_final, t_min_indices

class LIGHTS:
    def __init__(self, coord, color = None, intensity = 1): 
        self.coord = np.zeros((3))
        self.color = np.ones((3))
        self.intensity = intensity

        self.coord[:] = coord
        if color is not None:
            self.color[:] = color
        return

def is_collided_with_AABB(points, directions, box_center, box_size):
    """
    批量计算射线与AABB包围盒的相交情况
    参数:
        points: (N,3) 射线起点数组
        directions: (N,3) 射线方向数组(需归一化)
        box_center: (3,) 包围盒中心
        box_size: (3,) 包围盒尺寸(长宽高)
    返回:
        is_collided: (N,) bool数组，表示每条射线是否与包围盒相交
        t_min: (N,) 每条射线的最小相交参数
        t_max: (N,) 每条射线的最大相交参数
    """
    if points.ndim != 2 or directions.ndim != 2 or box_center.ndim != 1 or box_size.ndim != 1:
        raise ValueError('输入维度不正确')
    
    # 计算包围盒的最小/最大边界
    ray_origin = points
    ray_dir = directions
    box_min = box_center - box_size / 2
    box_max = box_center + box_size / 2
    
    # 初始化tmin和tmax
    t_min = np.full(ray_origin.shape[0], -np.inf)
    t_max = np.full(ray_origin.shape[0], np.inf)
    
    # 对每个维度进行处理
    for i in range(3):
        # 处理平行情况
        mask = np.abs(ray_dir[:, i]) < 1e-6
        outside = (ray_origin[:, i] < box_min[i]) | (ray_origin[:, i] > box_max[i])
        t_min[mask & outside] = np.inf
        t_max[mask & outside] = -np.inf
        
        # 非平行情况
        inv_dir = 1.0 / (ray_dir[:, i] + 1e-10)  # 避免除以零
        t1 = (box_min[i] - ray_origin[:, i]) * inv_dir
        t2 = (box_max[i] - ray_origin[:, i]) * inv_dir
        
        # 更新tmin和tmax
        t_min = np.maximum(t_min, np.minimum(t1, t2))
        t_max = np.minimum(t_max, np.maximum(t1, t2))
    
    # 相交判断
    is_collided = (t_max >= np.maximum(t_min, 0)) & (~np.isinf(t_max)) & (~np.isnan(t_max))
    return is_collided, t_min, t_max

@numba.njit(fastmath=True, cache=True)
def is_collided_with_AABB_numba(points, directions, box_center, box_size):
    """
    Numba加速版的射线与AABB包围盒相交检测
    参数:
        points: (N,3) 射线起点数组
        directions: (N,3) 射线方向数组(需归一化)
        box_center: (3,) 包围盒中心
        box_size: (3,) 包围盒尺寸(长宽高)
    返回:
        is_collided: (N,) bool数组，表示每条射线是否与包围盒相交
        t_min: (N,) 每条射线的最小相交参数
        t_max: (N,) 每条射线的最大相交参数
    """
    if points.ndim != 2 or directions.ndim !=2 or box_center.ndim != 1 or box_size.ndim != 1:
        raise ValueError('123')

    N = points.shape[0]
    box_min = box_center - box_size / 2
    box_max = box_center + box_size / 2
    
    t_min = np.full(N, -np.inf)
    t_max = np.full(N, np.inf)
    is_collided = np.full(N, False)
    
    for j in range(N):
        ray_origin = points[j]
        ray_dir = directions[j]
        current_t_min = -np.inf
        current_t_max = np.inf
        
        for i in range(3):
            # 处理平行情况
            if np.abs(ray_dir[i]) < 1e-6:
                if ray_origin[i] < box_min[i] or ray_origin[i] > box_max[i]:
                    current_t_min = np.inf
                    current_t_max = -np.inf
                    break
                continue
            
            # 非平行情况
            inv_dir = 1.0 / ray_dir[i]
            t1 = (box_min[i] - ray_origin[i]) * inv_dir
            t2 = (box_max[i] - ray_origin[i]) * inv_dir
            
            # 更新当前射线的tmin和tmax
            current_t_min = max(current_t_min, min(t1, t2))
            current_t_max = min(current_t_max, max(t1, t2))
        
        # 存储结果
        t_min[j] = current_t_min
        t_max[j] = current_t_max
        is_collided[j] = (current_t_max >= max(current_t_min, 0)) and (not np.isinf(current_t_max)) and (not np.isnan(current_t_max))
    
    return is_collided, t_min, t_max

def get_rotation_matrix(axis, theta):
    """
    返回绕任意轴旋转的3x3旋转矩阵

    参数:
        axis: 旋转轴向量 [kx, ky, kz]，需为单位向量
        theta: 旋转角度(弧度)

    返回:
        R: 3x3旋转矩阵
    """
    if axis.ndim != 1:
        raise ValueError('123')

    axis = axis / np.linalg.norm(axis)  # 归一化为单位向量
    kx, ky, kz = axis
    cos_theta = np.cos(theta)
    sin_theta = np.sin(theta)
    one_minus_cos = 1.0 - cos_theta

    R = np.array([
        [cos_theta + kx*kx*one_minus_cos, 
         kx*ky*one_minus_cos - kz*sin_theta, 
         kx*kz*one_minus_cos + ky*sin_theta],

        [ky*kx*one_minus_cos + kz*sin_theta, 
         cos_theta + ky*ky*one_minus_cos, 
         ky*kz*one_minus_cos - kx*sin_theta],

        [kz*kx*one_minus_cos - ky*sin_theta, 
         kz*ky*one_minus_cos + kx*sin_theta, 
         cos_theta + kz*kz*one_minus_cos]
    ])

    return R

def get_refelctive_direction(direction, norm):
    '''
    返回向量direction关于法向量norm的对称向量
    '''
    if direction.ndim != 2 or norm.ndim != 2:
        raise ValueError('123')

    direction = direction / np.linalg.norm(direction,axis = 1, keepdims = True)
    norm = norm / np.linalg.norm(norm,axis = 1, keepdims = True)

    def decompose_direction(direction, norm):
        """
        将 direction 向量分解为平行于 norm 和垂直于 norm 的两个部分。

        参数:
            direction: 向量，形状为 (3,) 或 (n,)
            norm: 单位法向量，形状与 direction 相同

        返回:
            parallel: 平行于 norm 的分量
            perpendicular: 垂直于 norm 的分量
        """

        # 计算点积
        dot_product = np.sum(direction * norm, axis = 1, keepdims = True)

        # 平行分量：投影到 norm 上
        parallel = dot_product * norm

        # 垂直分量：原向量减去平行分量
        perpendicular = direction - parallel

        return parallel, perpendicular
    parallel, perpendicular = decompose_direction(direction, norm)
    res = parallel - perpendicular
    res = res / np.linalg.norm(res, axis = 1, keepdims = True)
    return res

def load_texture(filename, target_res_height = 800):
    img = PIL.Image.open(filename).convert('RGB')
    width, height = img.size
    scale_factor = target_res_height / height
    new_width = int(width * scale_factor)
    new_height = int(height * scale_factor)
    img = img.resize((new_width, new_height), PIL.Image.LANCZOS)
    rgb_array = np.array(img, dtype=np.float64) / 255
    return np.clip(rgb_array, 0, 1)

def get_brightness(points, normals, view_directions, basecolor, metallic = 0.0, roughness = 0.5, depth = 0):
    if points.ndim != 2 or normals.ndim != 2 or view_directions.ndim != 2  or basecolor.ndim != 2:
        raise ValueError('123')

    N = points.shape[0]
    normals = normals / np.linalg.norm(normals, axis=1, keepdims=True)
    view_directions = - view_directions / np.linalg.norm(view_directions, axis=1, keepdims=True)

    L_tot = np.zeros((N,3))
    L_tot += ao * ambient_light_color[None,:] * basecolor #(1,3)

    for i in range(len(light_list)):
        light_ray = light_list[i].coord - points
        light_dist2 = np.sum(light_ray**2,axis = 1)
        light_ray = light_ray / np.linalg.norm(light_ray, axis=1, keepdims=True)
        light_color = light_list[i].color
        light_intensity = light_list[i].intensity * (1/(1+0.05*light_dist2)); # 半经验衰减。光照强度理论上按平方反比随距离衰减，但是单纯的平方反比衰减太快，又在近处容易发散，不容易控制

        is_collided = is_ray_collided_with_any_obj(points + 0.01 * light_ray, light_ray, use_Zbuffer=False)[0]

        # Compute half vector and NdotH
        H = (view_directions + light_ray) / 2.0
        H = H / np.linalg.norm(H, axis=1, keepdims=True)
        NdotL = np.clip(np.sum(normals * light_ray, axis=1), 1e-3, 1)
        NdotH = np.clip(np.sum(normals * H, axis=1), 1e-3, 1)

        # NdotL: 漫反射分量，基于朗伯余弦定律，与表面法向量和光源方向的夹角余弦成正比
        L_i = np.clip(1 - metallic*(1-roughness),0,1)[:,None] * light_intensity[:,None] * light_color[None, :] * basecolor * NdotL[:,None] #(N,3)

        # NdotH: 高光反射分量(Blinn-Phong模型)，当半角向量H接近法线时产生镜面高光
        L_i += light_intensity[:,None] * light_color[None, :] * basecolor * NdotL[:,None] * NdotH[:,None]**32 #(N,3)

        # Apply lighting only to non-collided points
        L_i[is_collided] = 0
        L_tot += L_i

    # 镜面采样
    if depth < 1:
        N_sample = 2
        reflective_directions = get_refelctive_direction(view_directions, normals) 
        _reflective_directions = reflective_directions[:,None,:] + 0.025*roughness[:,None,None]*(2*np.random.rand(N,N_sample,3) - 1) # 随机旋转
        B = np.array([start_raytracing(points + 0.01*_reflective_directions[:,i,:], _reflective_directions[:,i,:], depth = depth + 1)[0] for i in range(N_sample)]) # 递归采样 [2,N,3]
        reflective_intensity = np.clip(metallic*(1-roughness),0.025,1) # 反射强度
        LI = np.array([reflective_intensity[:,None] * B[i,:,:] * basecolor  for i in range(N_sample)]) #[2,N,3]
        L_i = np.mean(LI, axis=0)
        L_tot += L_i

    L_tot[np.isinf(L_tot) | np.isnan(L_tot)] = 0
    L_tot = np.clip(L_tot,0,1)
    return L_tot

def is_ray_collided_with_any_obj(points, directions, use_Zbuffer=True):
    N = points.shape[0]
    is_collided = np.full(N, False)
    collided_coord = np.empty((N, 3))
    normal = np.empty((N, 3))
    color = np.empty((N, 3))
    z_buffer = np.full(N, np.inf)
    collided_obj_id = np.full(N, -1, dtype=int)

    for i, obj in enumerate(obj_list):
        _is_collided, _collided_coord, _normal, _color, _t = obj.is_ray_collided_with_me(points, directions)

        update_mask = _is_collided & (~np.isinf(_t)) & (~np.isnan(_t)) & (_t > 0) & (_t < z_buffer)
        if np.any(update_mask):
            is_collided[update_mask] = True
            collided_coord[update_mask] = _collided_coord[update_mask]
            normal[update_mask] = _normal[update_mask]
            color[update_mask] = _color[update_mask]
            z_buffer[update_mask] = _t[update_mask]
            collided_obj_id[update_mask] = i

    return is_collided, collided_coord, normal, color, z_buffer, collided_obj_id

def start_raytracing(start_points, directions, depth = 0):
    N = start_points.shape[0]
    brightness = np.zeros((N, 3))

    # 检测碰撞
    is_collided, collided_coord, normal, color, t_obj, obj_id = is_ray_collided_with_any_obj(start_points, directions)

    # 计算亮度
    if np.any(is_collided):
        collided_mask = is_collided
        metallic = np.array([obj_list[i].advmat.metallic for i in obj_id[collided_mask]])
        roughness = np.array([obj_list[i].advmat.roughness for i in obj_id[collided_mask]])

        if is_pbr_ready:
            F0 = np.array([obj_list[i].advmat.F0 for i in obj_id[collided_mask]])

            b = drawsphere_pbr.get_brightness_pbr_simd(collided_coord[collided_mask], normal[collided_mask], directions[collided_mask], basecolors = color[collided_mask],
            F0 = F0,
            metallic = metallic,
            roughness = roughness,
            depth = depth
            )
        else:
            b = get_brightness(collided_coord[collided_mask], normal[collided_mask], directions[collided_mask], basecolor = color[collided_mask],metallic = metallic,roughness=roughness, depth = depth)
        brightness[collided_mask] = b

    return brightness, is_collided, t_obj, obj_id

ncam = 1.7

obj_list = []

#obj_list.append(SPHERES(np.array([0, 0, -5]), 1))
#'数据来源：Palmqvist, A. (2025) Physically Based - The PBR Values Database Available at: https://physicallybased.info/ (Accessed: 14 November 2025).'
#obj_list.append(SPHERES(np.array([1, 1, -4]), 0.5, color = np.array([0.910,0.778,0.423]), advmat = ADVMAT(F0 = np.array([0.995,0.974,0.747]), metallic = 1.0, roughness = 0.6)))
#obj_list.append(CUBES(np.array([0, -1.5, -9]), np.array([10, 1, 1]), color=np.array([0.0, 1.0, 0.0])))
#obj_list.append(SPRITERS(np.array([3, -2 + 1.5, -6]), 3, load_texture('image.jpg')))
#obj_list.append(GROUNDS(-2))

obj_triangles = TRIANGLES(coord = np.zeros((3)), model = scipy.io.loadmat('obj.mat')['triangles'], colors = np.array([0.910,0.778,0.423]), advmat = ADVMAT(F0 = np.array([0.995,0.974,0.747]), metallic = 1.0, roughness = 0.6))
obj_triangles.rotate(np.array([0,1,0]), 45/180*np.pi)
obj_triangles.scale(2)
obj_triangles.moveto(np.array([0,-1+obj_triangles.box_size[2]/2-0.5,-5]))

obj_list.append(SPHERES(np.array([-2, 1.5, -6]), 1))
obj_list.append(CUBES(np.array([0, -1.5, -5]), np.array([2.5,1,2.5]), color=np.array([0.0, 1.0, 0.0])))
obj_list.append(SPRITERS(np.array([3, -2 + 1.5, -6]), 3, load_texture('image.jpg')))
obj_list.append(obj_triangles)
obj_list.append(GROUNDS(-2))

light_list = []
light_list.append(LIGHTS(np.array([5,5,0]),color = np.array([1.0,1.0,1.0]),intensity = 4))

light_list = []
light_list.append(LIGHTS(np.array([5,5,0]),color = np.array([1.0,1.0,1.0]),intensity = 4))
ambient_light_color = np.array([1.0,1.0,1.0])
ao = 0.2

if is_pbr_ready:
    drawsphere_pbr.ao = ao
    drawsphere_pbr.ambient_light_color = ambient_light_color
    drawsphere_pbr.light_list = light_list
    drawsphere_pbr.is_ray_collided_with_any_obj_simd = is_ray_collided_with_any_obj
    drawsphere_pbr.start_raytracing_simd = start_raytracing
    drawsphere_pbr.get_refelctive_direction_simd = get_refelctive_direction


def apply_simple_fxaa(image, edge_threshold=0.125, edge_threshold_min=0.03125):
    from scipy.ndimage import gaussian_filter

    luma = np.dot(image[..., :3], [0.299, 0.587, 0.114])

    grad_x = np.zeros_like(luma)
    grad_y = np.zeros_like(luma)
    grad_x[:, 1:-1] = luma[:, 2:] - luma[:, :-2]
    grad_y[1:-1, :] = luma[2:, :] - luma[:-2, :]

    edge_intensity = np.abs(grad_x) + np.abs(grad_y)
    is_edge = edge_intensity > max(edge_threshold_min, edge_threshold * np.max(luma))

    blend_weights = np.where(is_edge, 0.5, 1.0)

    blurred = np.zeros_like(image)
    for c in range(3):
        blurred[..., c] = gaussian_filter(image[..., c], sigma=0.75)

    result = np.zeros_like(image)
    for c in range(3):
        result[..., c] = image[..., c] * blend_weights + blurred[..., c] * (1 - blend_weights)

    return np.clip(result, 0, 1)

def worker(task_queue, result_queue, ncam):
    """工作进程：从队列获取任务，计算后返回结果"""
    while True:
        task = task_queue.get()
        if task is None:  # 终止信号
            break

        idx_start, idx_end, batch_sxy = task
        batch_size = idx_end - idx_start

        # 创建射线起点和方向
        start_points = np.tile([0, 0, ncam], (batch_size, 1))
        directions = np.column_stack((batch_sxy, -ncam * np.ones(batch_size)))

        # 光线追踪计算
        brightness, is_collided, t_obj, obj_id = start_raytracing(start_points, directions)

        # 组装结果并发送回主进程
        batch_us = np.empty((batch_size, 5))
        batch_us[:, 0:3] = brightness
        batch_us[:, 3] = t_obj
        batch_us[:, 4] = obj_id

        result_queue.put((idx_start, idx_end, batch_us))

    return

if __name__ == '__main__':
    # 参数设置
    dx = 0.001
    L = 1

    # 生成输入数据
    xs, ys = np.meshgrid(np.arange(-L, L, dx), np.arange(-L, L, dx))
    sxy = np.column_stack((xs.flatten(), ys.flatten()))
    n = xs.shape[0]
    ns = n**2
    print(f"Total points: {ns}")

    # 创建队列
    task_queue = Queue()
    result_queue = Queue()

    start_time = time.time()

    # 启动工作进程
    num_proc = 6
    processes = []
    batch_size = 5000
    for _ in range(num_proc):
        p = Process(target=worker, args=(task_queue, result_queue, ncam))
        p.start()
        processes.append(p)

    # 分批发送任务
    for start_idx in range(0, ns, batch_size):
        end_idx = min(start_idx + batch_size, ns)
        batch_sxy = sxy[start_idx:end_idx]
        task_queue.put((start_idx, end_idx, batch_sxy))

    # 发送终止信号
    for _ in range(num_proc):
        task_queue.put(None)

    # 收集结果
    us = np.zeros((ns, 5))
    received_chunks = 0
    expected_chunks = (ns + batch_size - 1) // batch_size

    while received_chunks < expected_chunks:
        idx_start, idx_end, batch_us = result_queue.get()
        us[idx_start:idx_end] = batch_us
        received_chunks += 1
        print(f"Progress: {100*received_chunks/expected_chunks:.1f}%")

    # 等待所有进程结束
    for p in processes:
        p.join()

    print("Calculation complete!")
    print("Result shape:", us.shape)

    shared_us = us.copy()
    image = shared_us[:,0:3].reshape((n, n, 3))
    image = apply_simple_fxaa(image)
    image = (image * 255).astype(np.uint8)
    image = PIL.Image.fromarray(image).transpose(PIL.Image.FLIP_TOP_BOTTOM)

    image_deep = shared_us[:,3].reshape((n, n))
    image_deep[np.isnan(image_deep) | np.isinf(image_deep)] = 0
    image_deep = np.clip(image_deep,0,20) / 20
    image_deep = (image_deep * 255).astype(np.uint8)
    image_deep = PIL.Image.fromarray(image_deep).transpose(PIL.Image.FLIP_TOP_BOTTOM)

    image_objid = shared_us[:,4].reshape((n, n))
    image_objid = (image_objid+1)/(len(obj_list)+1)
    image_objid = (image_objid * 255).astype(np.uint8)
    image_objid = PIL.Image.fromarray(image_objid).transpose(PIL.Image.FLIP_TOP_BOTTOM)

    end_time = time.time()
    print(f"Execution time: {end_time - start_time:.2f} seconds")

    image.save('output_image.png')
    image_deep.save('output_image_deep.png')
    image_objid.save('output_image_objid.png')
    image.show()
